import os
import threading
import numpy as np
import networkx as nx
import rclpy
from rclpy.time import Time
from rclpy.duration import Duration
from rclpy.executors import ExternalShutdownException
from rclpy.executors import MultiThreadedExecutor
import tf2_ros
from semantic_nav.utils import SingleTaskExecutor
from semantic_nav.navigator import Navigator
from semantic_nav.topomap_builder import TopoMapBuilder
from semantic_nav.slam_utils import (pose_compose,
                                     pose_minus,
                                     pose_inv,
                                     Rt_to_PoseStamped,
                                     transform_to_Rt)
from semantic_nav.log_utils import get_logger
logger = get_logger()
package_name = 'semantic_nav'


class TopoMapNavigator(TopoMapBuilder):
    def __init__(self):
        super().__init__('topmap_navigator')
        self.navigator = Navigator()
        self.locator = None
        self.observations = None
        self.node_now = None
        self.time_stamp = None
        self.camera2base = None
        self.task_executor = SingleTaskExecutor()

    def init(self, path=None, slop=None):
        super().init(path=path, queue_size=5, slop=slop)

    def navigate(self, target):
        self.get_logger().info(f"navigating to {target}")
        while True:
            try:
                time_stamp, node_now, pose_in_node = self.task_executor.result()
                if pose_in_node is not None:
                    break
            except Exception as e:
                logger.error(f"Error finding pose in node: {e}")

        path = self.topomap.find_path(node_now, target, weight='dist')
        self.get_logger().info(f"path: {path}")
        final_pose = (np.eye(3), np.zeros(3))
        pose_in_node_camere = (pose_in_node[:3, :3], pose_in_node[:3, 3])
        pose_in_node_base = pose_compose(pose_compose(self.camera2base, pose_in_node_camere),
                                         pose_inv(self.camera2base))

        for i in range(len(path)-1):
            source_pose = self.topomap.node(path[i])['pose']
            target_pose = self.topomap.node(path[i+1])['pose']
            relative_pose = self.topomap.edge(path[i], path[i+1])['relative_pose']
            if path[i] < path[i+1]:
                final_pose = pose_compose(final_pose, relative_pose)
            else:
                final_pose = pose_compose(final_pose, pose_inv(relative_pose))

        final_pose = pose_minus(final_pose, pose_in_node_base)
        target_pose = Rt_to_PoseStamped(final_pose[0], final_pose[1], time_stamp, self.base_frame)
        self.navigator.navigate([target_pose])

    def process_observation(self, color_msgs, depth_msgs):
        while self.camera2base is None:
            try:
                tf = self.tf_buffer.lookup_transform(self.base_frame, color_msgs[0].header.frame_id, Time(), Duration(seconds=1.0))
                self.camera2base = transform_to_Rt(tf)
            except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
                self.get_logger().info(f"Error looking up transform: {e}")
        time_stamp, self.observations = self.convert_observation(color_msgs, depth_msgs)
        this_time = Time().from_msg(time_stamp)

        try:
            tf = self.tf_buffer.lookup_transform(self.map_frame, self.base_frame, this_time, Duration(seconds=1.0))
            this_pose = transform_to_Rt(tf)
            # node_now, relative_pose = self.topomap.locate(self.observations, self.camera_intrinsics, this_pose, time_stamp)
            self.task_executor.submit(self.topomap.locate, self.observations, self.camera_intrinsics, this_pose, time_stamp)
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            logger.error(f"Error looking up transform: {e}")


def main(args=None):
    rclpy.init(args=args)
    try:
        from ament_index_python.packages import get_package_prefix
        map_dir = os.path.join(get_package_prefix(package_name), 'maps')
        topomap_navigator = TopoMapNavigator()

        topomap_navigator.set_parameters([rclpy.parameter.Parameter(name="color_topics", value=["/zed/zed_node/rgb/image_rect_color"]),
                                          rclpy.parameter.Parameter(name="depth_topics", value=["/zed/zed_node/depth/depth_registered"]),
                                          rclpy.parameter.Parameter(name="color_info_topics", value=["/zed/zed_node/rgb/camera_info"]),
                                          rclpy.parameter.Parameter(name="depth_info_topics", value=["/zed/zed_node/depth/camera_info"])])

        topomap_navigator.init(map_dir)
        executor = MultiThreadedExecutor()
        executor.add_node(topomap_navigator)
        command_thread = threading.Thread(target=topomap_navigator.navigate, args=(0,))
        command_thread.start()
        executor.spin()
    except KeyboardInterrupt:
        pass
    except ExternalShutdownException:
        pass
    finally:
        command_thread.join()
        topomap_navigator.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
